import numpy as np
import tensorflow as tf
from tensorflow import keras
from PIL import Image
from tensorflow.keras import Sequential, layers

physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)

tf.random.set_seed(22)
np.random.seed(22)


# 把多张image保存达到一张image里面去。
def save_images(img, name):
    new_im = Image.new('L', (280, 280))

    index = 0
    for i in range(0, 280, 28):
        for j in range(0, 280, 28):
            im = img[index]
            im = Image.fromarray(im, mode='L')
            new_im.paste(im, (i, j))
            index += 1
    new_im.save(name)


# 定义超参数
h_dim = 20         # 把原来的784维护降低到20维度；
batchsise = 512      # fashion_mnist
lr = 1e-4

# 数据集加载
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train, x_test = x_train.astype(np.float32) / 255., x_test.astype(np.float32) / 255.


train_db = tf.data.Dataset.from_tensor_slices(x_train)
train_db = train_db.shuffle(batchsise * 5).batch(batchsise)

test_db = tf.data.Dataset.from_tensor_slices(x_test)
test_db = test_db.batch(100)


print(x_train.shape, y_train.shape)
print(x_test.shape, y_test.shape)


# 搭建模型
class AE(keras.Model):
    # 1. 初始化部分
    def __init__(self):
        super(AE, self).__init__()   # 调用父类的函数

        # Encoders编码, 网络
        self.encoder = Sequential([
            layers.Dense(256, activation=tf.nn.relu),
            layers.Dense(128, activation=tf.nn.relu),
            layers.Dense(h_dim)

        ])

        # Decoders解码，网络
        self.decoder = Sequential([
            layers.Dense(128, activation=tf.nn.relu),
            layers.Dense(256, activation=tf.nn.relu),
            layers.Dense(784)

        ])

    # 2. 前向传播的过程
    def call(self, inputs, training=None):
        # [b, 784] ==> [b, 10]
        h = self.encoder(inputs)
        # [b, 10] ==> [b, 784]
        x_ = self.decoder(h)

        return x_


# 创建模型
model = AE()
model.build(input_shape=(None, 784))
model.summary()
optimizer = keras.optimizers.Adam(lr=lr)

# 训练
h_random = tf.random.normal((100, 784)) # 指定随机编码，随后需要对随机编码重建，以了解自编码器对随机编码的复原能力
for epoch in range(50):
    for step, x in enumerate(train_db):
        # [b, 28, 28] => [b, 784]
        x = tf.reshape(x, [-1, 784]).numpy()
        with tf.GradientTape() as tape:
            x_rec_logits = model(x)

            # 把每个像素点当成一个二分类的问题；
            rec_loss = tf.losses.binary_crossentropy(x, x_rec_logits, from_logits=True)
            # rec_loss = tf.losses.MSE(x, x_rec_logits)
            rec_loss = tf.reduce_mean(rec_loss)

        grads = tape.gradient(rec_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        if step % 100 == 0:
            print('epoch: %3d, step:%4d, loss:%9f' % (epoch, step, float(rec_loss)))
            # evaluation
            x = next(iter(test_db))  # 从test_db中取出一张图片；
            x_shape = tf.reshape(x, [-1, 784]).numpy()
            logits = model(x_shape)  # 经过auto-encoder重建的效果。
            x_hat = tf.sigmoid(logits)  # 变化到0-1之间

            # [b, 784] => [b, 28, 28]       还原成原始尺寸；
            x_hat = tf.reshape(x_hat, [-1, 28, 28])  # 重建得到的图片；
            x_hat = x_hat.numpy() * 255.  # 把数据numpy取出来，变成0-255区间；
            x_hat = x_hat.astype(np.uint8)  # 还原成numpy保存图片的格式；
            save_images(x_hat, 'ae_image/rec_epoch_%d _step_%d.png' % (epoch, step))  # 每个epoch保存一次。

            h_random = tf.reshape(h_random, [-1, 28, 28])  # 重建得到的图片；
            h_random = h_random.numpy() * 255.  # 把数据numpy取出来，变成0-255区间；
            h_random = h_random.astype(np.uint8)  # 还原成numpy保存图片的格式；
            save_images(h_random, 'ae_image/rand_epoch_%d _step_%d.png' % (epoch, step))  # 每个epoch保存一次。
